"""训练+测试"""
import sys
sys.path.append('..\\energy')
sys.path.append('..\\force')
import torch
import datetime
from torch import nn
import torch.nn.functional as F
from cganet import GCNE
import numpy as np
from math import sqrt

#parameter
num_epochs = 300
batch_size = 64
learning_rate = 0.001
device_ids = [1] #要并行的GPU编号
torch.cuda.set_device(device_ids[0]) #主GPU

#固定随机数种子，使得结果可复现
"""np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True"""


#define
torch.set_default_dtype(torch.float32)
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
print(device)
model = GCNE(12,18,24,28)
model = nn.DataParallel(module=model, device_ids=device_ids)
model = model.cuda(device_ids[0]) #将模型放在device_ids[0]上作为controller
#model = torch.load('e4.61.6_Ag20_mse2.pth')
#model.to(device)
criterion = nn.MSELoss(reduction='mean')
#criterion = nn.L1Loss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
task_name = 'Ag34_12,18,24,28'

#load data
data = np.load('/home/hp/users/fuli/11_Ag_large/train/Ag34_CDEFGH3.npz')
temp = data['D']
spl = int(temp.shape[0]*0.9) #训练集和测试集划分
D = temp[:spl] #距离矩阵
D_test = temp[spl:]
E = data['E'][:spl] #能量
print(np.max(E))
E_test = data['E'][spl:]

#train
loss_file = open(task_name+'.txt','w')
loss_file.write(open('energy_train.py',encoding='utf-8').read()+'\n')
last_best_test = 1e10
for epoch in range(num_epochs):
    total_loss = 0
    for i in range(D.shape[0]//batch_size):
        d = torch.tensor(D[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        e = torch.tensor(E[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        outputs = model(d)
        loss = criterion(outputs, e)
        total_loss += sqrt(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if i%200==0:
            print ('  Batch [{}/{}], mean_loss={:.6f}'.format(i+1, D.shape[0]//batch_size, total_loss/(i+1)))
    
    info = 'Epoch [{:0>3d}/{}], mean loss={:.8f}'.format(epoch+1, num_epochs, total_loss/(i+1))
    print(info)
    loss_file.write(info)
    loss_file.flush()
    
    #test
    total_loss = 0
    for i in range(D_test.shape[0]//batch_size):
        d = torch.tensor(D_test[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        e = torch.tensor(E_test[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        outputs = model(d)
        loss = criterion(outputs, e)
        total_loss += sqrt(loss.item())
    info = ' '*4+'test error={:.8f}'.format(total_loss/(i+1))
    print(info)
    print(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f'))
    print(outputs[:20])
    print(e[:20])
    #save
    loss_file.write(info+'\n')
    loss_file.flush()
    if total_loss < last_best_test:
        last_best_test = total_loss
        torch.save(model,task_name+'.pth')
